#!/usr/bin/env python3
# ECMP + IPv4 P4Runtime controller (keeps original filename)
from p4runtime_lib.helper import P4InfoHelper
from p4runtime_lib.switch import SwitchConnection, ShutdownAllSwitchConnections

P4INFO = "build/load_balance.p4.p4info.txtpb"
BMV2JSON = "build/load_balance.json"

def program_pipeline(sw, p4info):
    sw.MasterArbitrationUpdate()
    sw.SetForwardingPipelineConfig(p4info=p4info.p4info, bmv2_json_file_path=BMV2JSON)

def add(sw, p4info, table, match, action, params):
    entry = p4info.buildTableEntry(table_name=table, match_fields=match,
                                   action_name=action, action_params=params)
    sw.WriteTableEntry(entry)

def main():
    p4info = P4InfoHelper(P4INFO)
    s = [SwitchConnection(f"s{i}", address=f"127.0.0.1:{50050+i}", device_id=i-1) for i in range(1,7)]

    try:
        for sw in s: program_pipeline(sw, p4info)

        # ---- s1 ----
        s1 = s[0]
        add(s1, p4info, "MyIngress.ecmp_group",
            {"hdr.ipv4.dstAddr":("10.0.0.1",32)},
            "MyIngress.set_ecmp_select", {"ecmp_base":0,"ecmp_count":4})

        for sel, (port,dmac) in enumerate([(2,"08:00:00:00:02:00"),
                                           (3,"08:00:00:00:03:00"),
                                           (4,"08:00:00:00:04:00"),
                                           (5,"08:00:00:00:05:00")]):
            add(s1, p4info, "MyIngress.ecmp_nhop",
                {"meta.ecmp_select": sel},
                "MyIngress.set_nhop", {"nhop_dmac": dmac, "nhop_ipv4":"10.0.0.1", "port": port})
        for port in (2,3,4,5):
            add(s1, p4info, "MyEgress.send_frame",
                {"standard_metadata.egress_port": port},
                "MyEgress.rewrite_mac", {"smac":"08:00:00:00:01:00"})

        add(s1, p4info, "MyIngress.ecmp_group",
            {"hdr.ipv4.dstAddr":("10.0.0.2",32)},
            "MyIngress.set_ecmp_select", {"ecmp_base":10,"ecmp_count":1})
        add(s1, p4info, "MyIngress.ecmp_nhop",
            {"meta.ecmp_select": 10},
            "MyIngress.set_nhop", {"nhop_dmac":"08:00:00:00:01:01","nhop_ipv4":"10.0.0.2","port":1})
        add(s1, p4info, "MyEgress.send_frame",
            {"standard_metadata.egress_port": 1},
            "MyEgress.rewrite_mac", {"smac":"08:00:00:00:01:00"})

        # ---- s6 ----
        s6 = s[5]
        add(s6, p4info, "MyIngress.ecmp_group",
            {"hdr.ipv4.dstAddr":("10.0.0.2",32)},
            "MyIngress.set_ecmp_select", {"ecmp_base":0,"ecmp_count":4})
        for sel, (port,dmac) in enumerate([(2,"08:00:00:00:02:00"),
                                           (3,"08:00:00:00:03:00"),
                                           (4,"08:00:00:00:04:00"),
                                           (5,"08:00:00:00:05:00")]):
            add(s6, p4info, "MyIngress.ecmp_nhop",
                {"meta.ecmp_select": sel},
                "MyIngress.set_nhop", {"nhop_dmac": dmac, "nhop_ipv4":"10.0.0.2", "port": port})
        for port in (2,3,4,5):
            add(s6, p4info, "MyEgress.send_frame",
                {"standard_metadata.egress_port": port},
                "MyEgress.rewrite_mac", {"smac":"08:00:00:00:06:00"})

        add(s6, p4info, "MyIngress.ecmp_group",
            {"hdr.ipv4.dstAddr":("10.0.0.1",32)},
            "MyIngress.set_ecmp_select", {"ecmp_base":10,"ecmp_count":1})
        add(s6, p4info, "MyIngress.ecmp_nhop",
            {"meta.ecmp_select": 10},
            "MyIngress.set_nhop", {"nhop_dmac":"08:00:00:00:02:02","nhop_ipv4":"10.0.0.1","port":1})
        add(s6, p4info, "MyEgress.send_frame",
            {"standard_metadata.egress_port": 1},
            "MyEgress.rewrite_mac", {"smac":"08:00:00:00:06:00"})

        # ---- s2..s5 transit ----
        for (idx, smac) in [(1,"08:00:00:00:02:00"),
                            (2,"08:00:00:00:03:00"),
                            (3,"08:00:00:00:04:00"),
                            (4,"08:00:00:00:05:00")]:
            sw = s[idx]
            add(sw, p4info, "MyIngress.ecmp_group",
                {"hdr.ipv4.dstAddr":("10.0.0.1",32)},
                "MyIngress.set_ecmp_select", {"ecmp_base":0,"ecmp_count":1})
            add(sw, p4info, "MyIngress.ecmp_nhop",
                {"meta.ecmp_select": 0},
                "MyIngress.set_nhop", {"nhop_dmac":"08:00:00:00:06:00","nhop_ipv4":"10.0.0.1","port":2})
            add(sw, p4info, "MyEgress.send_frame",
                {"standard_metadata.egress_port": 2},
                "MyEgress.rewrite_mac", {"smac":smac})

            add(sw, p4info, "MyIngress.ecmp_group",
                {"hdr.ipv4.dstAddr":("10.0.0.2",32)},
                "MyIngress.set_ecmp_select", {"ecmp_base":1,"ecmp_count":1})
            add(sw, p4info, "MyIngress.ecmp_nhop",
                {"meta.ecmp_select": 1},
                "MyIngress.set_nhop", {"nhop_dmac":"08:00:00:00:01:00","nhop_ipv4":"10.0.0.2","port":1})
            add(sw, p4info, "MyEgress.send_frame",
                {"standard_metadata.egress_port": 1},
                "MyEgress.rewrite_mac", {"smac":smac})

        print("ECMP + IPv4 rules installed. Press Enter to exit.")
        input()
    finally:
        ShutdownAllSwitchConnections()

if __name__ == "__main__":
    main()
